import torch
import os
from utils.oss_file import download_file
from models.model import ClassifyModel, SegmentModel, UnsupervisedModel
import uuid
from utils.mysql_crud import create_connection, select_data
from omegaconf import OmegaConf


class Export:
    def __init__(self, config):
        self.config = config

        if config.task_type == 'classify':
            self.model = ClassifyModel(config.classify_model.name, pretrained=False,
                                       class_num=config.classify_model.class_num, **config.classify_model.parameters)
            self.model = self.load_checkpoint_from_oss(self.model, config.classify_model.checkpoint_link,
                                                       config.classify_model.local_path)
        elif config.task_type == 'segment':
            self.model = SegmentModel(config.segment_model.name, pretrained=False,
                                      class_num=config.segment_model.class_num,
                                      **config.segment_model.parameters)
            self.model = self.load_checkpoint_from_oss(self.model, config.segment_model.checkpoint_link,
                                                       config.segment_model.local_path)

            if 'classify_model' in config:
                self.additional_model = ClassifyModel(config.classify_model.name, pretrained=False,
                                                      class_num=config.classify_model.class_num,
                                                      **config.classify_model.parameters)
                self.additional_model = self.load_checkpoint_from_oss(self.additional_model,
                                                                      config.classify_model.checkpoint_link,
                                                                      config.classify_model.local_path)

        elif config.task_type == 'unsupervised':
            self.model = UnsupervisedModel(config.unsupervised_model.name, **config.unsupervised_model.parameters)
            self.model = self.load_checkpoint_from_oss(self.model, config.unsupervised_model.checkpoint_link,
                                                       config.unsupervised_model.local_path)

        self.application_folder = f'./desktop_application/{uuid.uuid4()}'
        os.makedirs(self.application_folder)

    def load_checkpoint_from_oss(self, model, oss_link, local_folder):
        """
        http://msjava.oss-cn-qingdao.aliyuncs.com/train/log/2024/12/27/33582b55-92b1-4809-a3a5-06bd3d465ed6.pth
        """
        # 获取uuid
        if not os.path.exists(local_folder):
            os.makedirs(local_folder)

        checkpoint_name = oss_link.split('/')[-1]
        checkpoint_path = os.path.join(local_folder, checkpoint_name)
        if not os.path.isfile(checkpoint_path):
            print('download checkpoint from {}'.format(oss_link))
            download_file(oss_link, checkpoint_path)

        checkpoint = torch.load(checkpoint_path, map_location='cpu')
        model.load_state_dict(checkpoint['state_dict'])
        print(f'Successfully Loaded from {checkpoint_path}')
        model.eval()
        return model

    def export_model(self):
        input_names = ['input']
        output_names = ['output']
        batch_size = 1
        input_shape = (batch_size, 3, self.config.dataset.height, self.config.dataset.width)

        # if self.config.type == 'classify':
        #     output_shape = (batch_size, self.model.class_num)
        # elif self.config.type == 'segment':
        #     output_shape = (batch_size, self.model.class_num, self.config.dataset.height, self.config.dataset.width)
        # elif self.config.type == 'unsupervised':
        #     # todo
        #     pass

        torch.onnx.export(
            self.model,
            torch.randn(input_shape),
            self.application_folder + '/export.onnx',  # todo
            input_names=input_names,
            output_names=output_names,
            dynamic_axes={input_names[0]: {0: "batch_size"}, output_names[0]: {0: "batch_size"}}
            # 动态轴，即输入和输出张量可以具有不同的批次大小
        )
        print('Successfully exported to onnx')
        if hasattr(self, 'additional_model'):
            torch.onnx.export(
                self.additional_model,
                torch.randn(input_shape),
                self.application_folder + '/additional_export.onnx',  # todo
                input_names=input_names,
                output_names=output_names,
                dynamic_axes={input_names[0]: {0: "batch_size"}, output_names[0]: {0: "batch_size"}}
                # 动态轴，即输入和输出张量可以具有不同的批次大小
            )
            print('Successfully exported to additional onnx')

    def export_config2yaml(self):
        connection = create_connection()
        if self.config.task_type == 'classify' or self.config.task_type == 'segment':
            # 查询类别名保存至config
            query = f"""
            select a.name
                from dataset_category a, train_task b
            where b.id={self.config.train_task_id} and a.dataset_id = b.dataset_id and a.is_deleted = 0
            order by a.no
            """
            class_name = select_data(connection, query)
            class_name = [item[0] for item in class_name]
            self.config.dataset.class_name = class_name

        if self.config.task_type == 'segment':
            # 查询分割阈值
            query = f"SELECT other_res FROM train_task WHERE id = {self.config.train_task_id}"
            res = select_data(connection, query)
            res = [eval(item[0]) for item in res][0]
            self.config.best_thresholds = res['best_thresholds']
            self.config.best_minareas = res['best_minareas']
        OmegaConf.save(self.config, self.application_folder + '/config.yaml')
        connection.close()
